#include "LabelEquivalenceSeg.h"
#include <stdlib.h>
#include <iostream>

#include <cuda_runtime.h>

LabelEquivalenceSeg::LabelEquivalenceSeg(int width, int height):
show(height, width){
        this->width = width;
        this->height = height;

        cudaMalloc(&img, sizeof(n_label) * width * height);
        cudaMalloc(&ref, sizeof(int) * width * height);
        cudaMalloc(&counter, sizeof(int) * width * height);
        cudaMalloc(&boundary_device, sizeof(bool) * width * height);


        color_pool = new unsigned char*[width * height];
        boundary_host = new bool[width * height];
        label = new int[width * height];
        
        for(int i = 0; i < width * height; i++){
                color_pool[i] = new unsigned char[3];
                color_pool[i][0] = (unsigned char)rand();
                color_pool[i][1] = (unsigned char)rand();
                color_pool[i][2] = (unsigned char)rand();
        }
}


LabelEquivalenceSeg::~LabelEquivalenceSeg(){
        cudaFree(img);
        cudaFree(ref);
        cudaFree(counter);
        show.~Mat_<cv::Vec3b>();
        delete[] color_pool;
        delete[] boundary_host;
        delete[] label;
}

void LabelEquivalenceSeg::viewSegmentResult(){
        int* temp = new int[width * height];
        cudaMemcpy(temp, img, width * height * sizeof(int), cudaMemcpyDeviceToHost);
        cudaMemcpy(boundary_host, boundary_device, sizeof(bool) * width * height, cudaMemcpyDeviceToHost);
        

        cv::Vec3b* pixel;
        unsigned char* color;
        for(int i = 0; i < height; i++){
                for(int j = 0; j < width; j++){
                        pixel = &show.at<cv::Vec3b>(cv::Point2d(j, i));
                        if(temp[j + i * width] > -1 && boundary_host[j + i * width]){
                                color = color_pool[temp[j + i * width]];
                                pixel->val[0] = color[0];
                                pixel->val[1] = color[1];
                                pixel->val[2] = color[2];
                        } else {
                                pixel->val[0] = 0;
                                pixel->val[1] = 0;
                                pixel->val[2] = 0;
                        }
                }
        }

        delete[] temp;  
        cv::imshow("Labeled surfaces", show);
        cv::waitKey(1);
}

int* LabelEquivalenceSeg::getLabelMap() const{
        return label;
}

bool* LabelEquivalenceSeg::getBoundaryMap() const{
        return boundary_host;
}